#define _GNU_SOURCE
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <x86intrin.h>

#include "thp.h"
#include "netpbm.h"

#define EPSILON 0.005f

// Simulation parameters
static const unsigned int N = 500;

static const float SOURCE_TEMP = 5000.0f;
static const float BOUNDARY_TEMP = 2500.0f;

static const float MIN_DELTA = 0.05f;
static const unsigned int MAX_ITERATIONS = 10000;


static unsigned int idx(unsigned int x, unsigned int y, unsigned int stride) {
    return y * stride + x;
}


static void init(unsigned int source_x, unsigned int source_y, float * matrix) {
    // init
    memset(matrix, 0, N * N * sizeof(float));

    // place source
    matrix[idx(source_x, source_y, N)] = SOURCE_TEMP;

    // fill borders
    for (unsigned int x = 0; x < N; ++x) {
        matrix[idx(x, 0,   N)] = BOUNDARY_TEMP;
        matrix[idx(x, N-1, N)] = BOUNDARY_TEMP;
    }
    for (unsigned int y = 0; y < N; ++y) {
        matrix[idx(0,   y, N)] = BOUNDARY_TEMP;
        matrix[idx(N-1, y, N)] = BOUNDARY_TEMP;
    }
}

/*static void step(unsigned int source_x, unsigned int source_y, const float * current, float * next) {
    for (unsigned int y = 1; y < N-1; ++y) {
        for (unsigned int x = 1; x < N-1; ++x) {
            if ((y == source_y) && (x == source_x)) {
                continue;
            }
            next[idx(x, y, N)] = (current[idx(x, y-1, N)] +
                                  current[idx(x-1, y, N)] +
                                  current[idx(x+1, y, N)] +
                                  current[idx(x, y+1, N)]) / 4.0f;
        }
    }
}


static float diff(const float * current, const float * next) {
    float maxdiff = 0.0f;
    for (unsigned int y = 1; y < N-1; ++y) {
        for (unsigned int x = 1; x < N-1; ++x) {
            maxdiff = fmaxf(maxdiff, fabsf(next[idx(x, y, N)] - current[idx(x, y, N)]));
        }
    }
    return maxdiff;
}*/

static void step_sse(unsigned int source_x, unsigned int source_y, const float * current, float * next) {
    __m128 allfours = _mm_set1_ps(4.0f);

    for (unsigned int y = 1; y < N-1; ++y) {
        unsigned int x;

        for (x = 1; x <= (N-2) - (N-2) % 4; x+=4) {
            __m128 current4_prevcol = _mm_loadu_ps(&current[idx(x, y-1, N)]);
            __m128 current4_prevrow = _mm_loadu_ps(&current[idx(x-1, y, N)]);
            __m128 current4_nextrow = _mm_loadu_ps(&current[idx(x+1, y, N)]);
            __m128 current4_nextcol = _mm_loadu_ps(&current[idx(x, y+1, N)]);

            __m128 next4 = _mm_add_ps(current4_prevcol, current4_prevrow);
            next4 = _mm_add_ps(next4, current4_nextrow);
            next4 = _mm_add_ps(next4, current4_nextcol);

            next4 = _mm_div_ps(next4, allfours);

            _mm_storeu_ps(&next[idx(x, y, N)], next4);

       }

       next[idx(source_x, source_y, N)] = SOURCE_TEMP;

       for (; x < N-1; ++x) {
           if ((y == source_y) && (x == source_x)) {
               continue;
           }
           next[idx(x, y, N)] = (current[idx(x, y-1, N)] +
                                 current[idx(x-1, y, N)] +
                                 current[idx(x+1, y, N)] +
                                 current[idx(x, y+1, N)]) / 4.0f;
       }
    }
}


static float diff_sse(const float * current, const float * next) {
    __m128 maxdiff4 = _mm_set1_ps(0.0f);
    float maxdiff = 0.0f;

    for (unsigned int y = 1; y < N-1; ++y) {
        unsigned int x;

        for (x = 1; x <= (N-2) - (N-2) % 4; x+=4) {
            __m128 current4 = _mm_loadu_ps(&current[idx(x, y, N)]);
            __m128 next4 = _mm_loadu_ps(&next[idx(x, y, N)]);
            __m128 diff = _mm_sub_ps(next4, current4);
            __m128 ndiff = _mm_sub_ps(current4, next4);
            __m128 absdiff = _mm_max_ps(diff, ndiff);

            maxdiff4 = _mm_max_ps(maxdiff4, absdiff);
        }

        for(; x < N-1; ++x) {
            maxdiff = fmaxf(maxdiff, fabsf(next[idx(x, y, N)] - current[idx(x, y, N)]));
        }
    }

    maxdiff = fmaxf(maxdiff, fmaxf(maxdiff4[0], fmaxf(maxdiff4[1], fmaxf(maxdiff4[2], maxdiff4[3]))));

    return maxdiff;
}


int main() {
    size_t array_size = N * N * sizeof(float);

    /*float * current_gold = malloc(array_size);
    float * next_gold = malloc(array_size);*/
    float * current_sse = malloc_thp(array_size);
    float * next_sse = malloc_thp(array_size);

    srand(time(NULL));
    unsigned int source_x = rand() % (N-2) + 1;
    unsigned int source_y = rand() % (N-2) + 1;
    printf("Heat source at (%u, %u)\n", source_x, source_y);

    init(source_x, source_y, current_sse);
    /*memcpy(next_gold, current_gold, array_size);
    memcpy(current_sse, current_gold, array_size);*/
    memcpy(next_sse, current_sse, array_size);

    //float t_diff_gold = SOURCE_TEMP;
    float t_diff_sse = SOURCE_TEMP;
    for (unsigned int it = 0; (it < MAX_ITERATIONS) && (t_diff_sse > MIN_DELTA); ++it) {
        /*step(source_x, source_y, current_gold, next_gold);
        t_diff_gold = diff(current_gold, next_gold);
        printf("%u: %f\n", it, t_diff_gold);*/

        step_sse(source_x, source_y, current_sse, next_sse);
        t_diff_sse = diff_sse(current_sse, next_sse);
        printf("%u: %f\n", it, t_diff_sse);

        /*float * swap_gold = current_gold;
        current_gold = next_gold;
        next_gold = swap_gold;*/

        float * swap_sse = current_sse;
        current_sse = next_sse;
        next_sse = swap_sse;
    }

    /*int ok = 1;
    float max_epsilon = 0.0f;

    for (unsigned int y = 0; y < N; ++y) {
        for (unsigned int x = 0; x < N; ++x) {
            float ref = current_gold[idx(x, y, N)];
            float opt = current_sse[idx(x, y, N)];

            if (fabs(ref - opt) > max_epsilon) {
                max_epsilon = fabs(ref - opt);
            }

            if (fabs(ref - opt) > EPSILON) {
                printf("Mismatch at (%d, %d): %f vs %f\n", x, y, ref, opt);
                ok = 0;
                break;
            }
        }
    }

	if (ok) {
		printf("Todo ok\n");
	}

    if (max_epsilon > 0.0) {
        printf("Max epsilon: %.5f\n", max_epsilon);
    }*/

    char * filename;
    /*asprintf(&filename, "heat.gold.ppm");
    writePPMbinaryImage(filename, current_gold, N, N, N, 0.0f, fmaxf(SOURCE_TEMP, BOUNDARY_TEMP));*/
    asprintf(&filename, "heat.sse.ppm");
    writePPMbinaryImage(filename, current_sse, N, N, N, 0.0f, fmaxf(SOURCE_TEMP, BOUNDARY_TEMP));
    free(filename);

    /*free(current_gold);
    free(next_gold);*/
    free(current_sse);
    free(next_sse);

    return 0;
}
